from driver import lcm_subscriber, lcm_publisher
import lcm
import pybullet as p
import time
import pybullet_data
import random
from quad_bot import motors_info_pb

startPos = [0, 0, 0.25]
motorsIdList = [0, 1, 2, 4, 5, 6, 8, 9, 10, 12, 13, 14]
testTorques = [0, -1, 0, 0, -1, 0, 0, -1, 0, 0, -1, 0]
initialJointsPos = [0, -1, 2, 0, -1, 2, 0, -1, 2, 0, -1, 2]
pybulletTimeStep = 0.0005
count = 0
num = 0
_torques = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

# * Set pybullet environment
p.connect(p.GUI)
p.setGravity(0, 0, -9.81)
p.setTimeStep(pybulletTimeStep)
p.setAdditionalSearchPath(pybullet_data.getDataPath())  # optionally
planeId = p.loadURDF("plane.urdf")
quadrupedId = p.loadURDF("/urdf/mini_cheetah_me/mini_cheetah.urdf", startPos, useFixedBase=False)
# p.setRealTimeSimulation(1)
# print(p.getJointInfo(quadrupedId, 0))

# * Set motors to torque mode
for j in range(12):
    p.resetJointState(quadrupedId, motorsIdList[j], initialJointsPos[j])
for j in range(12):
    p.setJointMotorControl2(quadrupedId, motorsIdList[j], p.VELOCITY_CONTROL,force=0)
# print(p.getJointInfo(quadrupedId, 0))

# * Set sub and pub threads
sub = lcm_subscriber("test_subscriber", 0.0008, "MotorsTorquePB") #PB means PyBullet
pub = lcm_publisher("test_publisher", 0.001, "MotorsInfoPB","CheaterSensorInfoPB")
sub.start()
pub.start()

def updateMotorsInfo():
    joint_states = p.getJointStates(quadrupedId, motorsIdList)
    for i in range(12):
        pub.msg.pos[i] = joint_states[i][0]
        pub.msg.vel[i] = joint_states[i][1]


def updateIMUInfo():
    pos_ori = p.getBasePositionAndOrientation(quadrupedId)
    vel_omega = p.getBaseVelocity(quadrupedId) #World frame
    # print(pos_ori,vel_omega)
    for i in range(3):
        pub.msg2.pos[i] = pos_ori[0][i]
        pub.msg2.vel[i] = vel_omega[0][i]
        pub.msg2.omega[i] = vel_omega[1][i]
    for i in range(4):
        pub.msg2.quat[i] = pos_ori[1][i]

# p.resetBasePositionAndOrientation(quadrupedId,posObj=[1,0.5,0.25],ornObj=[ 0, 0, 0,1])

for i in range(10000000):

    # p.resetBaseVelocity(quadrupedId, linearVelocity=[0.2, 0, 0], angularVelocity=[0, 0, 0])
    p.stepSimulation()
    updateMotorsInfo()
    updateIMUInfo()
    # global count,num
    # if count>1000:
    #     count=1
    #     num+=1
    # print("subscriber run",num,count,"times")
    # count+=1
    # * PD control test
    # for i in range(4):
    #     _torques[i * 3] = 10 * (0 - pub.msg.pos[i * 3]) + 0.1*(0-pub.msg.vel[i*3])
    #     _torques[i * 3 + 1] = 10 * (-1 - pub.msg.pos[i * 3 + 1])+ 0.1*(0-pub.msg.vel[i*3+1])
    #     _torques[i * 3 + 2] = 10 * (2 - pub.msg.pos[i * 3 + 2])+ 0.1*(0-pub.msg.vel[i*3+2])
    p.setJointMotorControlArray(quadrupedId, motorsIdList, p.TORQUE_CONTROL, forces=sub.torques)
    # if (i < 10000):
    #     p.resetBasePositionAndOrientation(quadrupedId, posObj=[0., 0., 0.25],
    #                                       ornObj=[0, 0, 0., 1])  # [0, 0, 0.7071068, 0.7071068]
    # print(time.time())
    time.sleep(pybulletTimeStep)



sub.stop()
pub.stop()
p.disconnect()
